import cv2
import os
import sys

def extract_frames(video_path, output_dir, frame_interval=20):
    os.makedirs(output_dir, exist_ok=True) #如果目录已存在则不再创建；如果目录不存在则递归创建

    cap = cv2.VideoCapture(video_path)
    frame_idx = 0
    save_idx = 0

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        if frame_idx % frame_interval == 0:
            save_path = os.path.join(output_dir, f"frame_{save_idx}.jpg")
            cv2.imwrite(save_path, frame)
            save_idx += 1
        frame_idx += 1

    cap.release()
    print(f"提取了{save_idx}张帧")

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("用法: python collect_dataset.py <视频路径> <输出目录> [帧间隔]")
        sys.exit(1)
    
    video_path = sys.argv[1]
    output_dir = sys.argv[2]
    frame_interval = int(sys.argv[3]) if len(sys.argv) == 4 else 20
    extract_frames(video_path, output_dir, frame_interval)